require("dotenv").config();
const fs = require("fs");
const path = require("path");
const rp = require("request-promise");
const request = require("request");
const Sequelize = require("sequelize");
const { Consumer } = require("sqs-consumer");
const AWS = require("aws-sdk");
const streamingS3 = require("streaming-s3");
const archiver = require("archiver");
const to = require("await-to-js").default;
const rimraf = require("rimraf");
const POSTGRES_CONNECTION_STRING =
  process.env.POSTGRES_CONNECTION_STRING ||
  "postgres://postgres:password@localhost:6432/postgres";
const num_item_per_query = 10;
const num_item_per_progress_update = 100;
const bucket_name = "download.jilin-ai.com";
const ai_bucket_name = "ailib";
const http = require("https");
const uuidv1 = require("uuid/v1");

const downloadAsync = function(uri, filename) {
  return new Promise((resolve, reject) => {
    const file = fs.createWriteStream(filename);
    http.get(uri, function(res) {
      res.pipe(file);
      file.on("finish", function() {
        file.close();
        resolve(true);
      });
    });
  });
};

const awsConfig = {
  accessKeyId: "AKIAOUX2AD3PLYU5NO5Q",
  secretAccessKey: "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",
  region: "cn-northwest-1"
};
AWS.config.update(awsConfig);
const s3 = new AWS.S3();

const delay = timeout =>
  new Promise((resolve, reject) => {
    setTimeout(() => {
      resolve();
    }, timeout);
  });

const apAsync = function(options, name) {
  return new Promise((resolve, reject) => {
    var readStream = request(options);
    var uploader = new streamingS3(
      readStream,
      awsConfig,
      {
        Bucket: ai_bucket_name,
        Key: name,
        ACL: "public-read"
      },
      function(e, resp, stats) {
        if (e) {
          reject(e);
        } else {
          resolve(resp.Location);
        }
      }
    );
  });
};

const aiTaskWorker = async msg => {
  let queryParams, query, err, queryRes;
  const taskId = JSON.parse(msg.Body).id;
  try {
    var sequelize = new Sequelize(POSTGRES_CONNECTION_STRING, {});
    //get task
    [err, queryRes] = await to(
      sequelize.query("SELECT id, input from ai_lab_task WHERE id = :id", {
        replacements: {
          id: taskId
        }
      })
    );

    if (err) {
      console.log("set db task to working error: ", err);
      return;
    }

    //get image
    const inputObj = queryRes[0][0].input;
    const taskType = inputObj.type;
    const inputImageSrc = inputObj.url
      ? `result_${inputImageSrc}`
      : `${uuidv1}.jpg`;
    var options;
    //call ai api
    if (taskType === "maskrcnn") {
      options = {
        method: "POST",
        uri: "http://flower-boli.tpddns.cn:8080/api/test",
        form: {
          url: inputImageSrc
        }
      };
    } else if (taskType === "anime") {
      options = {
        method: "POST",
        uri: "http://flower-boli.tpddns.cn:8081/api/test"
      };
    }

    [err, queryRes] = await to(apAsync(options, inputImageSrc));

    if (err) {
      console.log("ai work error: ", err);
      return;
    }

    console.log("upload finish");

    [err, queryRes] = await to(
      sequelize.query(
        "UPDATE ai_lab_task SET state= :status, output= :output WHERE id = :id",
        {
          replacements: {
            id: taskId,
            status: "finished",
            output: JSON.stringify({
              url: queryRes
            })
          }
        }
      )
    );
    console.log("work ai task ok");
    return;
  } catch (e) {
    console.log(e);
    throw new Error(e);
  } finally {
    sequelize.close();
  }
};

const work = async msg => {
  const taskId = JSON.parse(msg.Body).id;
  const pid = JSON.parse(msg.Body).pid;
  let queryParams, query, err, queryRes;
  let totalItems = 0;
  let finishedItems = 0;
  //connect db
  var sequelize = new Sequelize(POSTGRES_CONNECTION_STRING, {});
  //set db task to working
  query = "UPDATE task SET status= :status WHERE id = :id";
  queryParams = {
    replacements: {
      id: taskId,
      status: "working"
    }
  };
  [err, queryRes] = await to(sequelize.query(query, queryParams));
  if (err) {
    console.log("set db task to working error: ", err);
    return;
  }
  //work
  //get finish item count
  query = "SELECT done from project_task_count_view where id = :id";
  queryParams = {
    replacements: {
      id: pid
    }
  };
  [err, queryRes] = await to(sequelize.query(query, queryParams));
  if (err) {
    console.log("set db task to working error: ", err);
    return;
  }

  if (queryRes[0].length > 0 || queryRes[0][0].done > 0) {
    let projectClasses;
    let projectDatasets;
    totalItems = queryRes[0][0].done;
    //get project classes and dataset servers
    query =
      "SELECT p.id,p.name from project_class as p where p.project_id = :pid";
    queryParams = {
      replacements: {
        pid: pid
      }
    };
    [err, queryRes] = await to(sequelize.query(query, queryParams));
    if (err) {
      console.log("get project classes  error: ", err);
      return;
    }
    if (!queryRes.length > 0 || !queryRes[0] > 0) {
      console.log("no project classes");
      return;
    }
    projectClasses = queryRes[0].reduce((pre, next) => {
      pre[next.id] = next.name;
      return pre;
    }, {});
    query =
      "select d.id,f.address " +
      "from project_dataset as pd,dataset as d,file_server as f " +
      "where pd.project_id = :pid " +
      "and pd.dataset_id = d.id " +
      "and f.id = d.file_server_id";
    queryParams = {
      replacements: {
        pid: pid
      }
    };
    [err, queryRes] = await to(sequelize.query(query, queryParams));
    if (err) {
      console.log("get dataset servers error: ", err);
      return;
    }
    if (!queryRes.length > 0 || !queryRes[0] > 0) {
      console.log("no project classes");
      return;
    }
    projectDatasets = queryRes[0].reduce((pre, next) => {
      pre[next.id] = next.address;
      return pre;
    }, {});

    //init folders
    let imgNameNum = 10000001;
    let folderNumber = Date.now();
    let folderURL = path.resolve(`./download/${folderNumber}`);
    let filesUrl = path.resolve(`./download/${folderNumber}/${folderNumber}`);
    let imgURL = path.resolve(
      `./download/${folderNumber}/${folderNumber}/images`
    );
    let labelURL = path.resolve(
      `./download/${folderNumber}/${folderNumber}/annotation`
    );
    let outputPath = path.join(`${folderURL}/${pid}.zip`);
    fs.mkdirSync(folderURL);
    fs.mkdirSync(filesUrl);
    fs.mkdirSync(imgURL);
    fs.mkdirSync(labelURL);
    //get items from db
    for (let i = 0; i < totalItems; i += num_item_per_query) {
      //get item with anns from db
      query =
        "SELECT p.id AS task_id,i.name As file_name,i.real_name as remote_name,i.dataset_id as dataset,json_agg(json_build_object('class', a.class_id,'conetent', a.content)) AS anns " +
        "FROM project_task_item AS p JOIN dataset_item AS i ON i.id = p.item_id JOIN annotation AS a ON a.task_item_id = p.id " +
        "WHERE p.project_id = :pid AND p.state = :state " +
        "GROUP BY p.id,i.name,i.real_name,i.dataset_id " +
        "LIMIT :limit OFFSET :offset";
      queryParams = {
        replacements: {
          pid: pid,
          state: "done",
          offset: i,
          limit: num_item_per_query
        }
      };
      [err, queryRes] = await to(sequelize.query(query, queryParams));
      if (err) {
        console.log("set db task to finished error: ", err);
        return;
      }
      if (queryRes.length > 0) {
        for (const item of queryRes[0]) {
          //download file from  remote
          const currImgNum = imgNameNum.toString().substring(1);
          const remoteUrl = `${projectDatasets[item.dataset]}/${
            item.remote_name
          }`;
          const fileType = item.file_name
            .substring(item.file_name.lastIndexOf("."), item.file_name.length)
            .toLowerCase();
          const detFile = `${imgURL}/${currImgNum}${fileType}`;
          [err, queryRes] = await to(downloadAsync(remoteUrl, detFile));
          if (err) {
            console.log("download image error", err);
          }
          //change class id to name
          item.anns.map((ann, i) => {
            const id = item.anns[i].class;
            item.anns[i].class = projectClasses[id];
          });
          // write labels to json file
          const detAnn = `${labelURL}/${currImgNum}.json`;
          fs.writeFileSync(detAnn, JSON.stringify(item));
          imgNameNum++;
        }
      }
    }
    const archive = archiver("zip");
    archive.directory(path.join(folderURL, "/"), false);
    let output = fs.createWriteStream(outputPath);

    output.on("close", function() {
      //upload to aws
      const fStream = fs.createReadStream(outputPath);
      var uploader = new streamingS3(
        fStream,
        awsConfig,
        {
          Bucket: bucket_name,
          Key: `${folderNumber}.zip`,
          ACL: "public-read"
        },
        async function(e, resp, stats) {
          if (e) return console.log("Upload error: ", e);
          //change task to finish and set zip file url and size
          query =
            "UPDATE task SET status= :status, output= :output WHERE id = :id";
          queryParams = {
            replacements: {
              id: taskId,
              output: JSON.stringify({
                url: resp.Location,
                size: stats.size
              }),
              status: "finished"
            }
          };
          [err, queryRes] = await to(sequelize.query(query, queryParams));
          if (err) {
            console.log("set db task to finished error: ", err);
            return;
          }
          //delete the folder
          rimraf.sync(folderURL);
        }
      );
    });

    archive.on("error", function(err) {
      console.log("archive error: ", err);
      throw err;
    });

    archive.pipe(output);
    archive.finalize();
  } else {
    query = "UPDATE task SET status= :status WHERE id = :id";
    queryParams = {
      replacements: {
        id: taskId,
        status: "finished"
      }
    };
    [err, queryRes] = await to(sequelize.query(query, queryParams));
    if (err) {
      console.log("set db task to finished error: ", err);
      return;
    }
  }
  //set db task to finished
};

const DownloadConsumer = Consumer.create({
  queueUrl:
    "https://sqs.cn-northwest-1.amazonaws.com.cn/535137518543/firsttest",
  handleMessage: async message => {
    await work(message);
  },
  sqs: new AWS.SQS()
});

DownloadConsumer.on("error", err => {
  console.error(err.message);
});

DownloadConsumer.on("processing_error", err => {
  console.error(err.message);
});

DownloadConsumer.on("timeout_error", err => {
  console.error(err.message);
});

DownloadConsumer.start();

const AiTaskConsumer = Consumer.create({
  queueUrl: "https://sqs.cn-northwest-1.amazonaws.com.cn/535137518543/aitask",
  handleMessage: async message => {
    await aiTaskWorker(message);
  },
  sqs: new AWS.SQS()
});

AiTaskConsumer.on("error", err => {
  console.error(err.message);
});

AiTaskConsumer.on("processing_error", err => {
  console.error(err.message);
});

AiTaskConsumer.on("timeout_error", err => {
  console.error(err.message);
});

AiTaskConsumer.start();
